package cc.jinglupeng.kom.pool;

import cc.jinglupeng.kom.bean.ZkBroker;
import cc.jinglupeng.kom.data.ZkData;
import kafka.javaapi.consumer.SimpleConsumer;
import org.apache.commons.pool.impl.GenericKeyedObjectPool;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.HashMap;
import java.util.Map;

/**
 * Created by jinglupeng on 4/10/16.
 */
public class KafkaConsumerPool extends GenericKeyedObjectPool<Integer, SimpleConsumer> {

    private static Logger logger = LoggerFactory.getLogger(KafkaConsumerPool.class);
    private static KafkaConsumerPool kafkaConsumerPool;
    private static Map<Integer, ZkBroker> zkBrokers = new HashMap<>();
    private static int defaultBrokerId;

    static {
        try {
            ZkBroker[] allBrokers = ZkData.getAllBrokers();
            if (allBrokers != null) {
                defaultBrokerId = allBrokers[0].getId();
                for (ZkBroker zkBroker : allBrokers) {
                    zkBrokers.put(zkBroker.getId(), zkBroker);
                }
            }
        } catch (Exception e) {
            logger.error("get all brokers failed", e);
        }
        KafkaConsumerFactory kafkaConsumerFactory = new KafkaConsumerFactory(zkBrokers);
        kafkaConsumerPool = new KafkaConsumerPool(kafkaConsumerFactory);
    }

    private KafkaConsumerPool(KafkaConsumerFactory kafkaConsumerFactory) {
        super(kafkaConsumerFactory);
    }


    public static SimpleConsumer borrowSimpleConsumer(Integer id) throws Exception {
        return kafkaConsumerPool.borrowObject(id);
    }

    public static void returnSimpleConsumer(Integer id, SimpleConsumer simpleConsumer) {
        try {
            kafkaConsumerPool.returnObject(id, simpleConsumer);
        } catch (Exception e) {
            logger.error("catch exception when return simpleConsumer", e);
        }
    }

    public static SimpleConsumer borrowDefaultSimpleConsumer() throws Exception {
        return borrowSimpleConsumer(defaultBrokerId);
    }

    public static void returnDefaultSimpleConsumer(SimpleConsumer simpleConsumer) {
        returnSimpleConsumer(defaultBrokerId, simpleConsumer);
    }
}
